Add quantization support for SHiRA#3321
Conversation
This builds upon PR huggingface#3117 and was used to review the effectiveness of the changes. Apparently the changes are good! On the way, some things were cleaned up like determining the in/out features using the appropriate utility function during init and masking.
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
|
@kkb-code We want to roll out quantization support to most PEFT methods. SHiRA is one of the first ones. If you have time, it would be great if you could double check. The idea is that with those relatively few changes, SHiRA will support multiple quantization types on the base model like BNB, torchao, and gptqmodel. |
BenjaminBossan
left a comment
BenjaminBossan
left a comment
There was a problem hiding this comment.
I gave this a try in the MetaMath benchmark and found two issues, please check.
| val = torch.ones_like(idx.type(base_layer.weight.dtype)) | ||
| mask = torch.zeros_like(base_layer.weight.view(1, -1)) | ||
| val = torch.ones(*idx.shape, dtype=bool) | ||
| mask = torch.zeros(*shape).to(val).view(1, -1) |
There was a problem hiding this comment.
I got an error when I ran this on CUDA because val and mask were on CPU. Something like this should work:
device = base_layer.weight.device
idx = (torch.randperm(base_layer.weight.numel(), generator=random_generator)[:num_shira_weights]).to(device)
val = torch.ones(*idx.shape, dtype=bool).to(device)
mask = torch.zeros(*shape).to(val).view(1, -1).to(device)
mask = mask.scatter_(1, idx.unsqueeze(0), val.unsqueeze(0)).view(shape)There was a problem hiding this comment.
I applied your suggestion.
| result = self.base_layer(x, *args, **kwargs) | ||
| else: | ||
| new_weight = copy.deepcopy(self.base_layer.weight.data) | ||
| base_result = self.base_layer(x) |
There was a problem hiding this comment.
With this change, we avoid the copy but we run the forward pass twice, once with the base layer and once with the SHiRA weights. On my machine, this runs considerably slower: train time 30s -> 80s, eval time 10s -> 70s.
There was a problem hiding this comment.
I just pushed changes that reverted back to the deepcopy variant (albeit using clone this time). Can you check again?
As suggested by Benjamin.
Apparently avoiding the copy is costing quite a lot of runtime performance. So, I reverted the change back to the iterative addition with final nn.linear call. The test and mask changes are necessary to make the random mask generation consistent across (and among) test calls. This was especially important for bnb 4bit quantization since `weight.numel()` for 4bit weights returns `numel // 2` due to packing.
BenjaminBossan
left a comment
BenjaminBossan
left a comment
There was a problem hiding this comment.
Thanks for the update. Unfortunately, I came across an issue when running the quantization tests with torchao_int8_dynamic_activation_int8. I described it below.
There are now two forward implementations: a slower fallback implementation using the base layer's forward and a faster implementation that uses the base weights instead, ignoring potential forward/backward hooks. Since the latter is significantly faster in training and inference it is a worthy implementation but we need at least warn the user that their forward/backward hooks are not being called.
a502c1f to
09bebe1
Compare
BenjaminBossan
left a comment
BenjaminBossan
left a comment
There was a problem hiding this comment.
Thanks for the update. Just two nits.
| return | ||
| self.scaling[adapter] = scale | ||
|
|
||
| @lru_cache(None) # noqa: B019 |
There was a problem hiding this comment.
For my understanding, we need the cache to prevent the warning to be fired by each affected layer? Let's maybe add a comment.
| output = peft_model(inputs) # should not raise | ||
| assert output.dtype == dtype | ||
|
|
||
| def test_shira_warns_about_hooks(self): |
There was a problem hiding this comment.
Can you explain why we expect hooks here?
There was a problem hiding this comment.
Done. I've also made the test more explicit and caught a bug in the caching (I wrongly assumed it was global).
The tests now make sure that the warning is only ever issued once and the code reflects that test. Previously, each instance of ShiraLayer had its own warning cache and therefore was able to emit the warning more than once. Now the state is global.
|
One of the new tests fails, could you please check? |
3 participants
This builds upon PR #3117 and was used to review the effectiveness of the changes. Apparently the changes are good!
On the way, some things were cleaned up like determining the in/out features using the appropriate utility function during init and masking.